import os

from PIL import Image, ImageDraw, ImageOps


def add_rounded_corners(image_path, radius, resolution_multiplier=4):
    """
    Add rounded corners to the given image.

    Args:
    image (PIL.Image): The input image.
    radius (int): The radius of the rounded corners.

    Returns:
    PIL.Image: The image with rounded corners.
    """
    # Create a mask with rounded corners
    img = Image.open(image_path).convert("RGBA")

    mask_size = (img.size[0] * resolution_multiplier, img.size[1] * resolution_multiplier)
    radius = resolution_multiplier * radius
    # aaa
    mask = Image.new("L", mask_size, 0)

    draw = ImageDraw.Draw(mask)

    draw.pieslice([0, 0, 2 * radius, 2 * radius], 180, 270, fill=255)
    draw.pieslice([0, mask_size[1] - 2 * radius, 2 * radius, mask_size[1]], 90, 180, fill=255)
    draw.pieslice([mask_size[0] - 2 * radius, mask_size[1] - 2 * radius, mask_size[0], mask_size[1]], 0, 90,
                  fill=255)
    draw.pieslice([mask_size[0] - 2 * radius, 0, mask_size[0], 2 * radius], 270, 0, fill=255)
    draw.rectangle([radius, 0, mask_size[0] - radius, mask_size[1]], fill=255)
    draw.rectangle([0, radius, mask_size[0], mask_size[1] - radius], fill=255)

    mask = mask.resize(img.size, Image.Resampling.LANCZOS)
    # Apply the mask to the image
    image_with_corners = ImageOps.fit(img, mask.size, centering=(0.5, 0.5))

    image_with_corners.putalpha(mask)

    return image_with_corners


def composite_on_black_background(image_with_corners, output_path, radius, border_width, resolution_multiplier=4):
    """
    Composite the given image with rounded corners on a black background.

    Args:
    image_with_corners (PIL.Image): The image with rounded corners.

    Returns:
    PIL.Image: The composite image.
    """
    # Create a black background mask
    size = (image_with_corners.size[0] + 2 * border_width, image_with_corners.size[1] + 2 * border_width)

    background = Image.new("RGBA", (size[0] * resolution_multiplier, size[1] * resolution_multiplier), (0, 0, 0, 0))
    draw = ImageDraw.Draw(background)

    # Draw rounded corners
    corner_radius = (radius + border_width) * resolution_multiplier
    # print(f'2 br is ${radius + border_width}')
    draw.pieslice([0, 0, corner_radius * 2, corner_radius * 2], 180, 270, fill=(0, 0, 0, 255))
    draw.pieslice([0, (size[1] * resolution_multiplier) - (corner_radius * 2), corner_radius * 2,
                   size[1] * resolution_multiplier], 90, 180, fill=(0, 0, 0, 255))
    draw.pieslice([(size[0] * resolution_multiplier) - (corner_radius * 2), 0, size[0] * resolution_multiplier,
                   corner_radius * 2], 270, 360, fill=(0, 0, 0, 255))
    draw.pieslice([(size[0] * resolution_multiplier) - (corner_radius * 2),
                   (size[1] * resolution_multiplier) - (corner_radius * 2), size[0] * resolution_multiplier,
                   size[1] * resolution_multiplier], 0, 90, fill=(0, 0, 0, 255))

    # Draw rectangles
    draw.rectangle(
        [corner_radius, 0, (size[0] * resolution_multiplier) - corner_radius, size[1] * resolution_multiplier],
        fill=(0, 0, 0, 255))
    draw.rectangle(
        [0, corner_radius, size[0] * resolution_multiplier, (size[1] * resolution_multiplier) - corner_radius],
        fill=(0, 0, 0, 255))

    background = background.resize(size, Image.Resampling.LANCZOS)

    # Composite the image on the black background
    background.paste(image_with_corners, (border_width, border_width), image_with_corners)
    # background.show()

    background.save(output_path, "PNG")


def process_images_in_directory(input_directory, output_directory):
    """
    处理输入目录中的所有图像并将结果保存到输出目录中。
    """
    # 确保输出目录存在
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    # 遍历输入目录中的所有文件
    for filename in os.listdir(input_directory):
        if filename.endswith(('.jpg', '.jpeg', '.png')):
            input_image_path = os.path.join(input_directory, filename)
            output_image_path = os.path.join(output_directory, filename)
            img = Image.open(input_image_path).convert("RGBA")
            radius = int(img.size[0] * (3 / 32))
            border_width = int(img.size[0] / 32)
            composite_on_black_background(add_rounded_corners(input_image_path, radius), output_image_path, radius,
                                          border_width)


'''
radius: w * (29/320)
border_width:w/32

'''

# Example usage:
input_path = r"C:\Users\Administrator\Desktop\test"
output_path = r"./demo"
process_images_in_directory(input_path, output_path)
